#!/usr/bin/python
"""
Author: Jimmy Saw
Usage: python dissertation_BLASTPLineageVotes.py blastfile.tbl 50
50 is max number of top hits to return
It is supposed to give you the parent phylum of the hit organism 
from a BLASTp result.
Needs to be in tabular format parsed by Blast2table Perl script
Last argument can take either phylum or class for now
Usage: python dissertation_BLASTPLineageVotes.py blastfile.blastp.tbl 20 phylum
Note: It's kind of slow because it queries NCBI Eutilities web services to
      retrieve top number of hits to calculate what is the top phylum.
"""

import sys
import re
from Bio import Entrez
from Bio import SeqIO

def fmt(f):
    st = '{0:.4}'.format(f)
    return st

Entrez.email = "jimmysaw@gmail.com"

blastfile = sys.argv[1]
bf = open(blastfile, "rU")
bfl = bf.readlines()

si = blastfile.split('.')
seqid = str(si[0])

maxorgs = sys.argv[2]

choice = str(sys.argv[3])

hit_acc = re.compile('\w+\|(\S+)\|')
orgname = re.compile('(\w+) (\w+)')

#p = seqid.match(blastfile)
#bfn = p.group(1)

num = len(bfl)

maximum = 0

if num > int(maxorgs):
    maximum = int(maxorgs)
else:
    maximum = num

phyla = []

if len(bfl) == 0:
    #print bfn + "\t" + "no BLAST hit"
    pass
else:
    for line in bfl[0:maximum]:
        s = line.split('\t')
        evalue = s[1]
        identities = s[3]
        similarities = s[4]
        locus_tag = s[5]
        accession = s[9]
        desc = s[13].rstrip()
        searchterm = ""
        taxo = ""
        pattern = hit_acc.match(accession)
        acc = pattern.group(1)
        handle = Entrez.efetch(db="nucleotide", id=acc, rettype="gb")
        rec = SeqIO.read(handle, "genbank")
        if choice == 'phylum':
            taxo = rec.annotations["taxonomy"][1]
        elif choice == 'class':
            taxo = rec.annotations["taxonomy"][2]
        elif choice == 'order':
            taxo = rec.annotations["taxonomy"][3]
        elif choice == 'genus':
            taxo = rec.annotations["taxonomy"][4]
        elif choice == 'species':
            print "Not implemented for species level"
        #print locus_tag + "\t" + taxo 
        phyla.append(taxo)

pls = len(phyla)

if pls > 0:
    unique_phyla = set(phyla)
    unique_phyla_list = list(unique_phyla)
    descriptions = []
    name_votes = []
    for d in unique_phyla_list:
        descriptions.append(d)
        count = phyla.count(d)
        name_votes.append(count)
    smallest = min(name_votes)
    largest = max(name_votes)
    max_index = name_votes.index(largest)
    min_index = name_votes.index(smallest)
    print seqid + "\t" + "Largest vote = " + str(largest) + "/" + str(maximum)\
        + " (" + str(fmt((float(largest)/float(maximum))*100)) + "%)" + "\t" \
        + descriptions[max_index]
    #print "Smallest vote = " + str(smallest) + "\t" + descriptions[min_index]
else:
    print seqid + "\t" + "No BLAST hit" + "\t" + "-"

bf.close()
